package raft

import (
	"bytes"
	"go-kvraft/labgob"
	"go-kvraft/rpcutil"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

type ApplyMsg struct {
	CommandValid bool
	Command      interface{}
	CommandIndex int

	// For 2D:
	SnapshotValid bool
	Snapshot      []byte
	SnapshotTerm  int
	SnapshotIndex int
}

// 枚举常量
const (
	Follower = iota
	Candidate
	Leader
	HeartBeatTimeout = 100
)

type Raft struct {
	mu        sync.RWMutex         // Lock to protect shared access to this peer's state
	peers     []*rpcutil.ClientEnd // RPC end points of all peers
	persister *Persister           // Object to hold this peer's persisted state
	me        int                  // this peer's index into peers[]
	dead      int32                // set by Kill()

	// Persistent state on all servers:
	// (Updated on stable storage before responding to RPCs)
	// 所有服务器上的持久性状态
	// (在响应RPC之前，在稳定的存储上更新)
	currentTerm int        // 服务器看到的最新任期（首次启动时初始化为0，单调增加）
	votedFor    int        // 当前任期内获得选票的候选人ID（如果没有则为空）
	log         []LogEntry // 日志条目；每个条目包含状态机的命令，以及领导者收到条目的时间（第一个索引是1）

	// Volatile state on all servers:
	// 所有服务器上的易失性状态
	commitIndex int // 已知被提交的最高日志条目的索引（初始化为0，单调增加）
	lastApplied int // 已应用于状态机的最高日志条目的索引（初始化为0，单调增加）

	// Volatile state on leaders:
	// (Reinitialized after election)
	// 领导者（服务器）上的易失性状态
	// (选举后重新初始化)
	nextIndex  []int // 对于每个服务器，要发送给该服务器的下一个日志条目的索引（初始化为领导者的最后一个日志索引+1）
	matchIndex []int // 对于每个服务器，已知在服务器上复制的最高日志条目的索引（初始化为0，单调增加）

	// 自己添加的一些：
	state         int
	lastHeartbeat time.Time // 最后一次心跳的时间戳
	voteCount     int       // 在某一任期内获得选票的数量

	applyCh   chan ApplyMsg // 用来发送应用信息给状态机的 channel
	applyCond *sync.Cond    //

	snapshot []byte
}

func (rf *Raft) GetState() (int, bool) {
	rf.mu.RLock()
	defer rf.mu.RUnlock()

	return rf.currentTerm, rf.state == Leader
}

// GetIndex 函数用于将日志索引转换为切片索引
func (rf *Raft) GetIndex(index int) int {
	return index - rf.log[0].Index
}

// GetLogEntry 函数用于获取指定索引的日志条目
func (rf *Raft) GetLogEntry(index int) LogEntry {
	if index < 0 {
		index += rf.log[0].Index + len(rf.log)
	}
	return rf.log[rf.GetIndex(index)]
}

func (rf *Raft) persist() {
	// Your code here (2C).
	writer := new(bytes.Buffer)
	encoder := labgob.NewEncoder(writer)
	encoder.Encode(rf.currentTerm) // 编码currentTerm
	encoder.Encode(rf.votedFor)    // 编码votedFor
	encoder.Encode(rf.log)         // 编码log
	data := writer.Bytes()
	rf.persister.SaveStateAndSnapshot(data, rf.snapshot) // 将编码后的数据保存到磁盘中
}

func (rf *Raft) readPersist(state []byte, snapshot []byte) {
	if state == nil || len(state) < 1 { // bootstrap without any state?
		return
	}
	// Your code here (2C).
	reader := bytes.NewBuffer(state)
	decoder := labgob.NewDecoder(reader)

	var currentTerm int
	var votedFor int
	var log []LogEntry
	// 解码currentTerm
	if decoder.Decode(&currentTerm) == nil {
		rf.currentTerm = currentTerm
	} else {
		return
	}
	// 解码votedFor
	if decoder.Decode(&votedFor) == nil {
		rf.votedFor = votedFor
	} else {
		return
	}

	// 解码log
	if decoder.Decode(&log) == nil {
		rf.log = log
		rf.lastApplied = rf.log[0].Index
		rf.commitIndex = rf.log[0].Index
	} else {
		return
	}
	rf.snapshot = snapshot
}

func (rf *Raft) Start(command interface{}) (int, int, bool) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	index := rf.GetLogEntry(-1).Index + 1
	term := rf.currentTerm
	isLeader := rf.state == Leader

	if isLeader {
		rf.log = append(rf.log, LogEntry{
			Command: command,
			Term:    term,
			Index:   index,
		})
		rf.matchIndex[rf.me] = index
		rf.nextIndex[rf.me] = index + 1
		rf.persist()

		rf.mu.Unlock()
		for server := range rf.peers {
			if server == rf.me {
				continue
			}
			go rf.leaderSendAppendEntries(server)
		}
		rf.mu.Lock()
	}
	return index, term, isLeader
}

func (rf *Raft) Kill() {
	atomic.StoreInt32(&rf.dead, 1)
	// Your code here, if desired.
}

func (rf *Raft) killed() bool {
	z := atomic.LoadInt32(&rf.dead)
	return z == 1
}

func (rf *Raft) electionRoutine() {
	for rf.killed() == false {

		// Your code here (2A)
		// Check if a leader election should be started.

		// pause for a random amount of time between 50 and 350
		// milliseconds.
		// ms := 50 + (rand.Int63() % 300)
		// time.Sleep(time.Duration(ms) * time.Millisecond)

		electionTimeout := time.Duration(HeartBeatTimeout*2+rand.Intn(HeartBeatTimeout)) * time.Millisecond
		rf.mu.Lock()
		// 如果当前服务器不是Leader，并且距离上次收到心跳消息的时间超过了选举超时时间，则开始选举
		if rf.state != Leader && time.Since(rf.lastHeartbeat) >= electionTimeout {
			// DPrintf("electionRoutine(): server[%d] election timeout, become candidate and start leader election", rf.me)
			rf.state = Candidate
			rf.currentTerm += 1
			rf.voteCount = 1
			rf.votedFor = rf.me
			rf.lastHeartbeat = time.Now()
			rf.persist()

			rf.mu.Unlock()

			// 向其他节点发送 RequestVote 请求
			for server := range rf.peers {
				if server == rf.me {
					continue
				}
				// DPrintf("electionRoutine(): candidate[%d] sendRequestVote to server[%d]", rf.me, server)
				go rf.candidateSendRequestVote(server)
			}
		} else {
			rf.mu.Unlock()
		}
		time.Sleep(10 * time.Millisecond)
	}
}

func (rf *Raft) replicationRoutine() {
	for rf.killed() == false {
		rf.mu.RLock()
		if rf.state == Leader {
			// 如果服务器是 Leader，则向其他服务器发起 AppendEntries 或者是 InstallSnapshot 请求
			for server := range rf.peers {
				if server == rf.me {
					continue
				}
				if rf.nextIndex[server] <= rf.log[0].Index {
					// 如果当前服务器的下一个需要接收的 LogEntry 的index 已经不在内存中了，直接甩一个 snapshot 过去给他
					rf.mu.RUnlock()
					go rf.sendInstallSnapshot(server)
				} else {
					// 否则发送 AppendEntries，心跳
					rf.mu.RUnlock()
					go rf.leaderSendAppendEntries(server)
				}
				rf.mu.RLock()
			}
		}
		rf.mu.RUnlock()
		time.Sleep(HeartBeatTimeout * time.Millisecond)
	}
}

func (rf *Raft) applyRoutine() {
	for rf.killed() == false {
		rf.mu.Lock()
		// 如果 lstApplied >= commitIndex，则等待
		for rf.lastApplied >= rf.commitIndex {
			rf.applyCond.Wait()
		}
		commitIndex := rf.commitIndex
		queue := make([]ApplyMsg, 0)
		// 将从lastApplied+1到commitIndex的日志条目应用到状态机
		for i := rf.lastApplied + 1; i <= rf.commitIndex; i++ {
			queue = append(queue, ApplyMsg{
				CommandValid: true,
				Command:      rf.GetLogEntry(i).Command,
				CommandIndex: rf.GetLogEntry(i).Index,
			})
		}
		rf.mu.Unlock()
		for _, applyMsg := range queue {
			rf.applyCh <- applyMsg
		}

		rf.mu.Lock()

		rf.lastApplied = Max(rf.lastApplied, commitIndex)
		rf.mu.Unlock()
	}
}

func Make(peers []*rpcutil.ClientEnd, me int, persister *Persister, applyCh chan ApplyMsg) *Raft {
	rf := &Raft{}
	rf.peers = peers
	rf.persister = persister
	rf.me = me

	// Your initialization code here (2A, 2B, 2C).
	rand.Seed(time.Now().UnixNano())

	// 初始化日志
	rf.log = make([]LogEntry, 0)
	rf.log = append(rf.log, LogEntry{0, 0, 0})

	rf.snapshot = make([]byte, 0)

	// 初始化状态
	rf.state = Follower
	rf.currentTerm = 0
	rf.votedFor = -1
	rf.lastHeartbeat = time.Now()

	// 初始化应用通道
	rf.applyCh = applyCh
	rf.applyCond = sync.NewCond(&rf.mu)

	// 初始化 index
	rf.commitIndex = 0
	rf.lastApplied = 0
	rf.matchIndex = make([]int, len(peers))
	rf.nextIndex = make([]int, len(peers))

	// initialize from state persisted before a crash
	rf.readPersist(persister.ReadRaftState(), persister.ReadSnapshot())

	// start electionRoutine goroutine to start elections
	go rf.electionRoutine()
	go rf.replicationRoutine()
	go rf.applyRoutine()
	return rf
}
